package com.weijian.memory;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.weijian.common.MessageWrapper;
import com.weijian.dal.entity.BaseConversation;
import com.weijian.service.MpConversationService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 消息记忆和数据库之间的关联
 */
@Service
@RequiredArgsConstructor
public class DbChatMemory implements ChatMemory {

    private final MpConversationService conversationService;


    /**
     *添加新的对话记录到数据库
     */
    @Override
    public void add(String conversationId, Message message) {
        add(conversationId, List.of(message));
    }

    /**
     *添加新的对话集合记录到数据库
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        BaseConversation baseConversation = conversationService.getOneOpt(
                        new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId))
                .orElseThrow(
                        () -> new RuntimeException(String.format("can not find the conversation , conversation id is %s", conversationId)));
        List<Message> savedMessages = convert(baseConversation);
        savedMessages.addAll(messages);
        conversationService.update(BaseConversation.builder()
                        .content(MessageWrapper.toConversationStr(savedMessages.stream().map(MessageWrapper::new).toList())).build(),
                new LambdaQueryWrapper<BaseConversation>()
                        .eq(BaseConversation::getConversationId, conversationId));
    }

    /**
     *获取指定conversationId下最新的lastN条消息
     */
    @Override
    public List<Message> get(String conversationId) {
        return conversationService.getOneOpt(new LambdaQueryWrapper<BaseConversation>()
                        .eq(BaseConversation::getConversationId, conversationId))
                .map(v -> convert(v).stream().toList())
                .orElse(List.of());
    }

    /**
     *清空历史对话
     */
    @Override
    public void clear(String conversationId) {
        conversationService.remove(
                new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId));
    }

    /**
     *将数据库中保存的对话记录转换为Message对象
     */
    private List<Message> convert(BaseConversation conversation) {
        if (!StringUtils.hasText(conversation.getContent())) {
            return new ArrayList<>();
        }
        return MessageWrapper.fromConversationStr(conversation.getContent()).stream().map(MessageWrapper::toMessage).collect(
                Collectors.toList());
    }
}